import asyncio
import struct
import threading
from datetime import datetime
from typing import Dict

from pymodbus import __version__ as pymodbus_version
from pymodbus.datastore import (
    ModbusSequentialDataBlock,
    ModbusServerContext,
    ModbusSlaveContext,
)
from pymodbus.device import ModbusDeviceIdentification
from pymodbus.server import (
    StartAsyncTcpServer,
    StartAsyncSerialServer,
    ServerStop,
)
from pymodbus.framer import ModbusRtuFramer

from service.common.message_queue import MessageQueue
from service.enums.decode import Decode, DecodeType
from service.log import log


class BaseModbusServer:
    def __init__(
            self,
    ):
        self.server = None
        self.loop = None
        self.thread = None
        self.message_queue = MessageQueue()
        self.stop_event = threading.Event()
        # 创建从站上下文
        self.slaves: Dict[int, ModbusSlaveContext] = {}
        self.context = ModbusServerContext(slaves=self.slaves, single=False)
        self.setUpServer()

    def request_tracer(self, request, *addr):
        """请求追踪回调函数"""
        pdu = bytes([request.function_code]) + request.encode()
        transaction_id = request.transaction_id
        protocol_id = 0x0000
        length = len(pdu) + 1
        unit_id = request.slave_id

        mbap_header = bytearray([
            (transaction_id >> 8) & 0xFF,
            transaction_id & 0xFF,
            (protocol_id >> 8) & 0xFF,
            protocol_id & 0xFF,
            (length >> 8) & 0xFF,
            length & 0xFF,
            unit_id,
        ])

        full_request = mbap_header + pdu
        request_hex = " ".join(f"{b:02x}" for b in full_request)
        receive_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
        receive_str = f"[{receive_time}] RX: {request_hex}"
        log.info(receive_str)
        self.message_queue.enqueue(receive_str)

    def response_manipulator(self, response):
        """响应处理回调函数"""
        pdu = bytes([response.function_code]) + response.encode()
        transaction_id = response.transaction_id
        protocol_id = 0x0000
        length = len(pdu) + 1
        unit_id = response.slave_id

        mbap_header = bytearray([
            (transaction_id >> 8) & 0xFF,
            transaction_id & 0xFF,
            (protocol_id >> 8) & 0xFF,
            protocol_id & 0xFF,
            (length >> 8) & 0xFF,
            length & 0xFF,
            unit_id,
        ])

        full_response = mbap_header + pdu
        response_hex = " ".join(f"{b:02x}" for b in full_response)
        send_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
        send_str = f"[{send_time}] TX: {response_hex}"
        log.info(send_str)
        self.message_queue.enqueue(send_str)
        return response, False

    def get_latest_message(self) -> list:
        """获取最新的消息"""
        message_list = []
        while not self.message_queue.is_empty():
            message_list.append(self.message_queue.dequeue())
        log.info(f"获取最新的消息: {message_list}")
        return message_list

    def is_slave_exist(self, slave_id: int) -> bool:
        """
        检查从站是否存在
        """
        return slave_id in self.slaves

    def add_slave(self, slave_id: int) -> bool:
        """
        添加从站
        """
        if slave_id not in self.slaves:
            self.slaves[slave_id] = ModbusSlaveContext(
                di=ModbusSequentialDataBlock(0, [0]),
                co=ModbusSequentialDataBlock(0, [0]),
                hr=ModbusSequentialDataBlock(0, [0]),
                ir=ModbusSequentialDataBlock(0, [0])
            )
            return True
        return False

    def delete_slave(self, slave_id: int) -> bool:
        """
        删除从站
        """
        if slave_id in self.slaves:
            del self.slaves[slave_id]
            return True
        return False

    def add_slave_di(self, slave_id: int, start_address: int, end_address: int):
        """
        外部调用该方法为指定从站添加离散输入寄存器地址范围。

        :param slave_id: 从站 ID
        :param di_range: 离散输入寄存器地址范围，元组格式 (起始地址, 结束地址)
        """
        di_block = ModbusSequentialDataBlock(start_address, [0] * (end_address - start_address + 1))
        if slave_id not in self.slaves:
            self.slaves[slave_id] = ModbusSlaveContext(
                di=di_block,
                co=ModbusSequentialDataBlock(0, [0]),  # 先初始化一个空的，后续可更新
                hr=ModbusSequentialDataBlock(0, [0]),
                ir=ModbusSequentialDataBlock(0, [0])
            )
        else:
            self.slaves[slave_id].store["d"] = di_block
        self.context = ModbusServerContext(slaves=self.slaves, single=False)

    def add_slave_co(self, slave_id: int, start_address: int, end_address: int):
        """
        外部调用该方法为指定从站添加线圈寄存器地址范围。

        :param slave_id: 从站 ID
        :param co_range: 线圈寄存器地址范围，元组格式 (起始地址, 结束地址)
        """
        co_block = ModbusSequentialDataBlock(start_address, [0] * (end_address - start_address + 1))
        if slave_id not in self.slaves:
            self.slaves[slave_id] = ModbusSlaveContext(
                di=ModbusSequentialDataBlock(0, [0]),
                co=co_block,
                hr=ModbusSequentialDataBlock(0, [0]),
                ir=ModbusSequentialDataBlock(0, [0])
            )
        else:
            self.slaves[slave_id].store["c"] = co_block
        self.context = ModbusServerContext(slaves=self.slaves, single=False)

    def add_slave_hr(self, slave_id: int, start_address: int, end_address: int):
        """
        外部调用该方法为指定从站添加保持寄存器地址范围。

        :param slave_id: 从站 ID
        :param hr_range: 保持寄存器地址范围，元组格式 (起始地址, 结束地址)
        """
        hr_block = ModbusSequentialDataBlock(start_address, [0] * (end_address - start_address + 1))
        if slave_id not in self.slaves:
            self.slaves[slave_id] = ModbusSlaveContext(
                di=ModbusSequentialDataBlock(0, [0]),
                co=ModbusSequentialDataBlock(0, [0]),
                hr=hr_block,
                ir=ModbusSequentialDataBlock(0, [0])
            )
        else:
            self.slaves[slave_id].store["h"] = hr_block
        self.context = ModbusServerContext(slaves=self.slaves, single=False)

    def add_slave_ir(self, slave_id: int, start_address: int, end_address: int):
        """
        外部调用该方法为指定从站添加输入寄存器地址范围。

        :param slave_id: 从站 ID
        :param ir_range: 输入寄存器地址范围，元组格式 (起始地址, 结束地址)
        """
        ir_block = ModbusSequentialDataBlock(start_address, [0] * (end_address - start_address + 1))
        if slave_id not in self.slaves:
            self.slaves[slave_id] = ModbusSlaveContext(
                di=ModbusSequentialDataBlock(0, [0]),
                co=ModbusSequentialDataBlock(0, [0]),
                hr=ModbusSequentialDataBlock(0, [0]),
                ir=ir_block
            )
        else:
            self.slaves[slave_id].store["i"] = ir_block
        self.context = ModbusServerContext(slaves=self.slaves, single=False)

    def setUpServer(self, description=None, context=None, cmdline=None):
        """Run server setup."""
        # 直接初始化数据存储
        log.info("### Create datastore")
        # 固定使用 sequential 模式
        datablock = ModbusSequentialDataBlock(0x00, [17] * 100)

        if context is None:
            # 假设使用多个从站
            context = {
                0x01: ModbusSlaveContext(
                    di=datablock,
                    co=datablock,
                    hr=datablock,
                    ir=datablock,
                ),
                0x02: ModbusSlaveContext(
                    di=datablock,
                    co=datablock,
                    hr=datablock,
                    ir=datablock,
                ),
                0x03: ModbusSlaveContext(
                    di=datablock,
                    co=datablock,
                    hr=datablock,
                    ir=datablock,
                    zero_mode=True,
                ),
            }
            single = False
        else:
            single = True

        server_context = ModbusServerContext(slaves=context, single=single)

        # 初始化设备标识
        self.identity = ModbusDeviceIdentification(
            info_name={
                "VendorName": "Pymodbus",
                "ProductCode": "PM",
                "VendorUrl": "https://github.com/pymodbus-dev/pymodbus/",
                "ProductName": "Pymodbus Server",
                "ModelName": "Pymodbus Server",
                "MajorMinorRevision": pymodbus_version,
            }
        )

        return {
            "context": server_context,
            "identity": self.identity,
        }

    async def runAsyncServer(self):
        raise NotImplementedError

    async def initServer(self):
        runParams = self.setUpServer(
            description="Run callback server.", cmdline=None, context=self.context
        )
        self.server = await self.runAsyncServer()

    def startAsyncServer(self):
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        self.loop.run_until_complete(self.initServer())
        self.loop.close()

    async def _run_server(self):
        """异步运行服务器初始化方法。"""
        await self.initServer()

    def startServer(self):
        """直接在主线程启动异步服务器。"""
        self.stop_event.clear()
        try:
            self.loop = asyncio.get_running_loop()
        except RuntimeError:
            self.loop = asyncio.new_event_loop()
            asyncio.set_event_loop(self.loop)
        self.loop.create_task(self._run_server())
        if not self.loop.is_running():
            self.loop.run_forever()

    def stopServer(self):
        """停止异步服务器。"""
        log.info("停止运行modbus_server")
        if self.loop is not None:
            if self.loop.is_running():
                self.loop.call_soon_threadsafe(ServerStop)
                self.loop.call_soon_threadsafe(self.loop.stop)
                log.info("停止运行modbus_server协程")
            self.stop_event.set()
        self.loop = None

    def isRunning(self):
        return self.thread is not None and self.thread.is_alive()

    def setValueByAddress(
            self,
            func_code,
            rtu_addr,
            address,
            value,
            decode="0x20",
            register_cnt=1,
            signed=False,
    ):
        func_code = int(func_code)
        rtu_addr = int(rtu_addr)
        decode_type = Decode.get_decode_type(decode)
        endian = Decode.get_endian(decode)

        if register_cnt == 2:
            if decode_type == DecodeType.Float:
                packed = struct.pack(f"{endian}f", float(value))
                registers = list(struct.unpack(f"{endian}HH", packed))
            else:
                fmt = f"{endian}l" if signed else f"{endian}L"
                packed = struct.pack(fmt, int(value))
                registers = list(struct.unpack(f"{endian}HH", packed))
        elif register_cnt == 1:
            value = int(value)
            if signed and value < 0:
                value = (1 << 16) + value
            registers = [value & 0xFFFF]
            if endian == "<":
                registers[0] = ((registers[0] & 0xFF) << 8) | ((registers[0] >> 8) & 0xFF)

        if func_code == 10:
            func_code = 6
        self.slaves[rtu_addr].setValues(func_code, address, registers)

    def getValueByAddress(
            self,
            func_code,
            rtu_addr,
            address,
            decode="0x20",
            register_cnt=1,
            signed=False,
    ):
        func_code = int(func_code)
        rtu_addr = int(rtu_addr)
        if func_code == 10:
            func_code = 6

        raw_values = self.slaves[rtu_addr].getValues(func_code, address, register_cnt)
        if not raw_values:
            return 0

        endian = Decode.get_endian(decode)
        if register_cnt == 2:
            packed = struct.pack(f"{endian}HH", *raw_values)
            if Decode.get_decode_type(decode) == DecodeType.Float:
                return struct.unpack(f"{endian}f", packed)[0]
            else:
                fmt = f"{endian}l" if signed else f"{endian}L"
                return struct.unpack(fmt, packed)[0]
        else:
            value = raw_values[0]
            if endian == "<":
                value = ((value & 0xFF) << 8) | ((value >> 8) & 0xFF)
            if signed and value > 0x7FFF:
                value -= 0x10000
            return value


class ModbusTcpServer(BaseModbusServer):
    def __init__(
            self,
            port: int = 502,
            server_address: str = "0.0.0.0"
    ):
        super().__init__()
        self.port = port
        self.server_address = server_address

    def setServerAddress(self, address):
        self.server_address = address

    async def runAsyncServer(self):
        log.info(f"### start ASYNC TCP server, listening on {self.port}")
        address = (self.server_address, self.port)
        return await StartAsyncTcpServer(
            context=self.context,
            identity=self.identity,
            address=address,
            request_tracer=self.request_tracer,
            response_manipulator=self.response_manipulator
        )


class ModbusSerialServer(BaseModbusServer):
    def __init__(
            self,
            serial_port: str = "/dev/ttyUSB0",
            stopbits: int = 1,
            bytesize: int = 8,
            parity: str = "N",
            baudrate: int = 9600,
    ):
        super().__init__()
        self.serial_port = serial_port
        self.stopbits = stopbits
        self.bytesize = bytesize
        self.parity = parity
        self.baudrate = baudrate

    def setSerialPort(self, serial_port):
        self.serial_port = serial_port

    async def runAsyncServer(self):
        log.info(f"### start ASYNC Serial server, listening on {self.serial_port}")
        return await StartAsyncSerialServer(
            context=self.context,
            identity=self.identity,
            port=self.serial_port,
            stopbits=self.stopbits,
            bytesize=self.bytesize,
            parity=self.parity,
            baudrate=self.baudrate,
            request_tracer=self.request_tracer,
            response_manipulator=self.response_manipulator
        )


class ModbusRtuOverTcpServer(BaseModbusServer):
    def __init__(
            self,
            port: int = 502,
            server_address: str = "0.0.0.0"
    ):
        super().__init__()
        self.port = port
        self.server_address = server_address

    async def runAsyncServer(self):
        log.info(f"### start ASYNC RTU over TCP server, listening on {self.port}")
        return await StartAsyncTcpServer(
            context=self.context,
            identity=self.identity,
            port=self.port,
            framer=ModbusRtuFramer,
            server_address=self.server_address,
            request_tracer=self.request_tracer,
            response_manipulator=self.response_manipulator
        )
